from django.conf import settings
from tasks.models import DemoTask
from tasks.helper import get_loop_redis
from django.core.management import BaseCommand
from tasks.helper import friend_loads_json, friend_dumps_json

import asyncio, json, time
JobRegister = {
    DemoTask.job_name: DemoTask
}



class Command(BaseCommand):
    """ 添加任务 """

    async def add_task(self, redis):
        while True:
            key_job = await redis.blpop(settings.REDIS_TASK_ADD, timeout=10)
            if not key_job or key_job[0].decode() != settings.REDIS_TASK_ADD: continue
            print(key_job)
            jobs, obj_lst = friend_loads_json(key_job[1]), list()
            for job in jobs:
                expire_at = int(time.time()) + DemoTask.ExpireLimit * len(jobs)
                job['expire_at'] = expire_at
                if JobRegister.get(job.get('name', '')):
                    obj_lst.append(job)
            if obj_lst:
                for obj in await JobRegister.get(job.get('name', '')).insert_many(obj_lst):
                    asyncio.ensure_future(
                        redis.rpush(settings.REDIS_TASK_TAG,
                                    friend_dumps_json({'job_name': obj.name, 'job_id': obj.job_hash})))

    def cancel_task(self):
        for task in asyncio.Task.all_tasks():
            task.cancel()

    def handle(self, *args, **options):
        loop, redis = get_loop_redis()
        loop.run_until_complete(self.add_task(redis))
        self.cancel_task()

# class Command(BaseCommand):
#
#     async def add_task(self, redis):
#         test_tasks = ['https://www.pianshen.com/article/32141247147/'] * 10
#         while True:
#             expire_at = int(time.time()) + DemoTask.ExpireLimit * len(test_tasks)
#             tasks = [
#                 {'url': url, 'expire_at': expire_at, 'name': DemoTask.job_name}
#                 for url in test_tasks]
#             task_objs = await DemoTask.insert_many(tasks)
#             for obj in task_objs:
#                 job = {'job_name': DemoTask.job_name, 'job_id': obj.job_hash}
#                 asyncio.ensure_future(redis.rpush(settings.REDIS_TASK_TAG, json.dumps(job)))
#             await asyncio.sleep(2)
#
#     def cancel_task(self):
#         for task in asyncio.Task.all_tasks():
#             task.cancel()
#
#     def handle(self, *args, **options):
#         loop, redis = get_loop_redis()
#         loop.run_until_complete(self.add_task(redis))
#         self.cancel_task()
